# -*- coding: utf-8 -*-
import jsonpickle
from celery import Celery
from celery._state import _task_stack
from flask import Flask, has_request_context, request, g, has_app_context
from celery import Task


def celery_init_app(app: Flask) -> Celery:
    """
    创建celery对象
    :param app:
    :return:
    """

    class FlaskTask(Task):

        CONTEXT_ARG_NAME = '_flask_request_context'

        def __call__(self, *args, **kwargs) -> object:
            _task_stack.push(self)
            self.push_request(args=args, kwargs=kwargs)
            try:
                with app.app_context():
                    callback = lambda: self.run(*args, **kwargs)

                    context = kwargs.pop(self.CONTEXT_ARG_NAME, None)

                    flask_g = kwargs.pop('__flask_g', '{}')

                    for k, v in jsonpickle.loads(flask_g).items():
                        setattr(g, k, v)

                    if context is None or has_request_context():
                        return callback()

                    with app.test_request_context(**context):
                        return callback()
            finally:
                self.pop_request()
                _task_stack.pop()

        def apply_async(self, args=None, kwargs=None, **rest):
            self._load_request_context(kwargs)
            self._load_app_context(kwargs)
            # arguments  = inspect.signature(super(FlaskTask, self).apply_async).bind(*args, **kwargs).arguments.copy()
            return super(FlaskTask, self).apply_async(args, kwargs, **rest)

        def apply(self, args=None, kwargs=None, **rest):
            self._load_request_context(kwargs)

            return super(FlaskTask, self).apply(args, kwargs, **rest)

        def retry(self, args=None, kwargs=None, **rest):
            self._load_request_context(kwargs)

            return super(FlaskTask, self).retry(args, kwargs, **rest)

        def _load_request_context(self, kwargs):
            """加载Flask请求上下文"""
            if not has_request_context():
                return

            # keys correspond to arguments of :meth:`Flask.test_request_context`
            context = {
                'path': request.path,
                'base_url': request.url_root,
                'method': request.method,
                'headers': dict(request.headers),
                'data': request.form or (request.is_json) and request.json or {}
            }
            if '?' in request.url:
                context['query_string'] = request.url[(request.url.find('?') + 1):]

            if kwargs is None:
                kwargs = {}

            kwargs[self.CONTEXT_ARG_NAME] = context

        def _load_app_context(self, kwargs):
            """加载Flask请求上下文"""
            if not has_app_context():
                return
            if kwargs is None:
                kwargs = {}
            kwargs['__flask_g'] = jsonpickle.dumps(g.__dict__)

        def on_success(self, retval, task_id, args, kwargs):
            print(f'task done: {retval}')
            return super(FlaskTask, self).on_success(retval, task_id, args, kwargs)

        def on_failure(self, exc, task_id, args, kwargs, einfo):
            print(f'task fail, reason: {exc}')
            return super(FlaskTask, self).on_failure(exc, task_id, args, kwargs, einfo)

    celery_app = Celery(app.name, task_cls=FlaskTask)
    celery_app.config_from_object(app.config["CELERY"])
    celery_app.set_default()
    app.extensions["celery"] = celery_app
    return celery_app
